package com.hy.aicodegeneration.core.pool;

import org.openqa.selenium.*;
import org.openqa.selenium.interactions.Coordinates;
import org.openqa.selenium.interactions.Locatable;

import java.util.List;
import java.util.Set;

/**
 * 池化的 WebDriver 包装类
 */
public class PooledWebDriver implements WebDriver, JavascriptExecutor, TakesScreenshot, SearchContext {

    private final WebDriver webDriver;
    private final WebDriverPool pool;
    private boolean borrowed = true;

    public PooledWebDriver(WebDriver webDriver, WebDriverPool pool) {
        this.webDriver = webDriver;
        this.pool = pool;
    }

    /**
     * 获取原始 WebDriver 实例
     *
     * @return 原始 WebDriver
     */
    public WebDriver getWebDriver() {
        return webDriver;
    }

    /**
     * 归还连接到池中
     */
    public void returnToPool() {
        if (borrowed) {
            pool.returnDriver(this);
            borrowed = false;
        }
    }

    // ==================== WebDriver 接口实现 ====================

    @Override
    public void get(String url) {
        webDriver.get(url);
    }

    @Override
    public String getCurrentUrl() {
        return webDriver.getCurrentUrl();
    }

    @Override
    public String getTitle() {
        return webDriver.getTitle();
    }

    @Override
    public List<WebElement> findElements(By by) {
        return webDriver.findElements(by);
    }

    @Override
    public WebElement findElement(By by) {
        return webDriver.findElement(by);
    }

    @Override
    public String getPageSource() {
        return webDriver.getPageSource();
    }

    @Override
    public void close() {
        webDriver.close();
    }

    @Override
    public void quit() {
        webDriver.quit();
    }

    @Override
    public Set<String> getWindowHandles() {
        return webDriver.getWindowHandles();
    }

    @Override
    public String getWindowHandle() {
        return webDriver.getWindowHandle();
    }

    @Override
    public TargetLocator switchTo() {
        return webDriver.switchTo();
    }

    @Override
    public Navigation navigate() {
        return webDriver.navigate();
    }

    @Override
    public Options manage() {
        return webDriver.manage();
    }

    // ==================== JavascriptExecutor 接口实现 ====================

    @Override
    public Object executeScript(String script, Object... args) {
        return ((JavascriptExecutor) webDriver).executeScript(script, args);
    }

    @Override
    public Object executeAsyncScript(String script, Object... args) {
        return ((JavascriptExecutor) webDriver).executeAsyncScript(script, args);
    }

    // ==================== TakesScreenshot 接口实现 ====================

    @Override
    public <X> X getScreenshotAs(OutputType<X> target) throws WebDriverException {
        return ((TakesScreenshot) webDriver).getScreenshotAs(target);
    }
}
